import os
import cv2
import datetime
import shutil
import numpy as np
from PIL import ImageFont, ImageDraw, Image
class DatasetCreator:
    def __init__(self):
        self.root_dir = ""
        self.train_dir = ""
        self.test_dir = ""
        self.test_dir = ""
        self.train_dir_class_name_file = ""
        self.val_dir_class_name_file = ""
        self.test_dir_class_name_file = ""        
        self.pic_count = 1  
        self.train_count = 1  # 新增
        self.val_count = 1  # 新增
        self.test_count = 1  # 新增  
        self.class_flags = {}  # 新增
    def create_file(self,root_dir,class_name):
        try:
            self.root_dir = root_dir
            self.train_dir = os.path.join(root_dir, "training_set")
            self.val_dir = os.path.join(root_dir, "val_set")
            self.test_dir = os.path.join(root_dir, "test_set")
            self.train_dir_class_name_file = os.path.join(self.train_dir, class_name)
            self.val_dir_class_name_file = os.path.join(self.val_dir, class_name)
            self.test_dir_class_name_file = os.path.join(self.test_dir, class_name)
            # 创建目录，如果它们不存在
            os.makedirs(self.train_dir, exist_ok=True)
            os.makedirs(self.val_dir, exist_ok=True)
            os.makedirs(self.test_dir, exist_ok=True)
            os.makedirs(self.train_dir_class_name_file, exist_ok=True)
            os.makedirs(self.val_dir_class_name_file, exist_ok=True)
            os.makedirs(self.test_dir_class_name_file, exist_ok=True)
            # 检查并追加关键词到 classes.txt

            classes_file = os.path.join(self.root_dir, "classes.txt")
            if os.path.exists(classes_file):
                with open(classes_file, 'r') as f:
                    classes = f.read().splitlines()
                if class_name in classes:
                    # print(f"Keyword '{class_name}' already exists in classes.txt.")
                    return
            else:
                classes = []

            with open(classes_file, 'a') as f:
                f.write(f"{class_name}\n")
            # 在创建新类别的文件夹时，重置该类别的标志
            self.class_flags[class_name] = False
            self.pic_count = 1
            self.train_count = 1
            self.val_count = 1
            self.test_count = 1
        except Exception as e:
            print(f"Failed to create directories for category '{class_name}'. Reason: {e}")
    def clear_category_images(self, root_dir,class_name):
        # 清空指定分类的图片
        try:
            self.root_dir = root_dir
            self.train_dir = os.path.join(root_dir, "training_set")
            self.val_dir = os.path.join(root_dir, "val_set")
            self.test_dir = os.path.join(root_dir, "test_set")
            train_class_dir = os.path.join(self.train_dir, class_name)
            val_class_dir = os.path.join(self.val_dir, class_name)
            test_class_dir = os.path.join(self.test_dir, class_name)
            # 清空训练集中的分类图片
            self.clear_images(train_class_dir)
            # 清空验证集中的分类图片
            self.clear_images(val_class_dir)
            # 清空测试集中的分类图片
            self.clear_images(test_class_dir)
        except Exception as e:
            print(f"Failed to clear images for category '{class_name}'. Reason: {e}")

    def clear_images(self, directory):
        try:
            if os.path.exists(directory):
                # 删除目录下的所有文件，但保留子目录
                for filename in os.listdir(directory):
                    file_path = os.path.join(directory, filename)
                    try:
                        if os.path.isfile(file_path) or os.path.islink(file_path):
                            os.unlink(file_path)
                        # 如果需要打印被删除的文件名，您可以在这里添加打印语句
                        # print(f'Deleted {file_path}')
                    except Exception as e:
                        print(f'Failed to delete {file_path}. Reason: {e}')
                # print(f"Images in '{directory}' have been cleared.")

            # else:
            #     print(f"The directory {directory} does not exist.")
        except Exception as e:
            print(f"Failed to clear images in '{directory}'. Reason: {e}") 
    def clear_category(self, root_dir, class_name):
        # 清空指定分类的目录
        try:
            self.root_dir = root_dir
            self.train_dir = os.path.join(root_dir, "training_set")
            self.val_dir = os.path.join(root_dir, "val_set")
            self.test_dir = os.path.join(root_dir, "test_set")
            train_class_dir = os.path.join(self.train_dir, class_name)
            val_class_dir = os.path.join(self.val_dir, class_name)
            test_class_dir = os.path.join(self.test_dir, class_name)
  

            # 删除训练集中的分类目录
            if os.path.exists(train_class_dir):
                shutil.rmtree(train_class_dir)
            # 删除验证集中的分类目录
            if os.path.exists(val_class_dir):
                shutil.rmtree(val_class_dir)
            # 删除测试集中的分类目录
            if os.path.exists(test_class_dir):
                shutil.rmtree(test_class_dir)
            classes_file = os.path.join(self.root_dir, "classes.txt")
            class_exists = False
            # 检查类别名称是否存在于classes.txt中，并记录下来
            if os.path.exists(classes_file):
                with open(classes_file, 'r') as f:
                    classes = f.read().splitlines()
                class_exists = class_name in classes
            # else:
            #     print(f"The file {classes_file} does not exist.")
            # 如果存在，删除类别名称
            if class_exists:
                # 删除关键词
                classes.remove(class_name)
                with open(classes_file, 'w') as f:
                    for cls in classes:
                        f.write(f"{cls}\n")
                # print(f"Keyword '{class_name}' has been removed from classes.txt.")
            # else:
            #     print(f"Keyword '{class_name}' does not exist in classes.txt.")
        except Exception as e:
            print(f"Failed to clear category '{class_name}' directories. Reason: {e}")

    def display_image(self, cap,img_src_x,img_src_y):
        cvimg_success, img_src = cap.read()
        if not cvimg_success:
            return
        h, w, c = img_src.shape
        start_x = (w-h)//2
        start_y = (h-h)//2
        end_x = start_x+h
        end_y = start_y+h
        img_src = img_src[start_y:end_y, start_x:end_x]
        img_src_a = cv2.resize(img_src,(img_src_x,img_src_y))
        img_src = self.drawChinese(text=str(self.pic_count)+str(".jpg"),x=10, y=10,size=20,r=255,g=0,b=0,a=0,img=img_src_a)
        cv2.imshow('cvwindow', img_src)
        cv2.waitKey(10)
        return img_src_a
    
    def drawChinese(self,text,x,y,size,r, g, b, a,img):
        font = ImageFont.truetype("HYQiHei_50S.ttf", size)
        img_pil = Image.fromarray(img)
        draw = ImageDraw.Draw(img_pil)
        draw.text((x,y), text, font=font, fill=(b, g, r, a))
        frame = np.array(img_pil)
        return frame
    def collect_and_sort_images(self,class_name,img_src, total_num, train_ratio, val_ratio, test_ratio):
            if self.class_flags.get(class_name, False):
                print(f"All images of class '{class_name}' have been collected.")
                return
            epsilon = 0.01
            assert abs((train_ratio + val_ratio + test_ratio) - 1) < epsilon, "Ratios must sum up to 1"
            print(total_num,self.pic_count)
            train_num = int(total_num * train_ratio)
            val_num = int(total_num * val_ratio)
            test_num = total_num - train_num - val_num

            if self.pic_count <= total_num:

                img_name_path = str(self.pic_count) + ".jpg"
                
                # Save the image to the appropriate directory based on the count
                if self.test_count <= test_num:
                    save_path = os.path.join(self.test_dir_class_name_file, img_name_path)
                    self.test_count += 1
                elif self.train_count <= train_num:
                    save_path = os.path.join(self.train_dir_class_name_file, img_name_path)
                    self.train_count += 1
                elif self.val_count <= val_num:
                    save_path = os.path.join(self.val_dir_class_name_file, img_name_path)
                    self.val_count += 1
                else:
                    # 如果所有的条件都不满足，就抛出一个异常
                    raise RuntimeError("Unable to determine the save path.")
                
                cv2.imwrite(save_path, img_src)
                print(str(self.pic_count)+str(".jpg"))
                self.pic_count += 1
            else:
                print(f"All images of class '{class_name}' have been collected.")
                self.pic_count = 1
                self.train_count = 0
                self.val_count = 0
                self.test_count = 0
                self.class_flags[class_name] = True 
            cv2.waitKey(10)

